import os
import pywt
import json
import pickle
import argparse

import numpy as np
import pandas as pd
from collections import OrderedDict
from sklearn.preprocessing import MinMaxScaler, StandardScaler
from sklearn.decomposition import KernelPCA

from pyhht import EMD, plot_imfs
from scipy.signal import hilbert
import matplotlib.pyplot as plt

PHM_RUL = {'Bearing1_1': 0, 'Bearing1_2': 0,
           'Bearing2_1': 0, 'Bearing2_2': 0,
           'Bearing3_1': 0, 'Bearing3_2': 0,
           'Bearing1_3': 5730, 'Bearing1_4': 339, 'Bearing1_5': 1610, 'Bearing1_6': 1460, 'Bearing1_7': 7570,
           'Bearing2_3': 7530, 'Bearing2_4': 1390, 'Bearing2_5': 3090, 'Bearing2_6': 1290, 'Bearing2_7': 580,
           'Bearing3_3': 820}
PHM_TOKEN = {'Bearing1_1': 0, 'Bearing1_2': 0,
             'Bearing2_1': 1, 'Bearing2_2': 1,
             'Bearing3_1': 2, 'Bearing3_2': 2,
             'Bearing1_3': 0, 'Bearing1_4': 0, 'Bearing1_5': 0, 'Bearing1_6': 0, 'Bearing1_7': 0,
             'Bearing2_3': 1, 'Bearing2_4': 1, 'Bearing2_5': 1, 'Bearing2_6': 1, 'Bearing2_7': 1,
             'Bearing3_3': 2}

Recession_point = {'Bearing1_1': 27400, 'Bearing1_2': 8270, 'Bearing1_3': 22430, 'Bearing1_4': 10870,
                   'Bearing1_5': 24100,
                   'Bearing1_6': 24140, 'Bearing1_7': 22110}

PHM_LIFE = 18000
# PHM_LIFE = 1
# TRAIN_BEARING = ["Bearing1_1", "Bearing1_2", "Bearing2_1", "Bearing2_2", "Bearing3_1", "Bearing3_2"]

# TEST_BEARING = ["Bearing1_3", "Bearing1_4",  "Bearing1_5", "Bearing1_6", "Bearing1_7", "Bearing2_3",
#                 "Bearing2_4", "Bearing2_5", "Bearing2_6", "Bearing2_7", "Bearing3_3"]

# TRAIN_BEARING = ["Bearing3_1", "Bearing3_2"]
# TEST_BEARING = ["Bearing3_3"]


TRAIN_BEARING = ["Bearing1_1", "Bearing1_2", "Bearing1_4", "Bearing1_5", "Bearing1_6", "Bearing1_7"]
TEST_BEARING = ["Bearing1_3"]


# TRAIN_BEARING = ["Bearing2_1", "Bearing2_2"]
# TEST_BEARING = ["Bearing2_3", "Bearing2_4",  "Bearing2_5", "Bearing2_6", "Bearing2_7"]
def main(data_path, save_path, need_process, full_test):
    os.makedirs(save_path, exist_ok=True)

    # 加载数据并保存
    data_dict = load_phm2012(data_path, need_process, True, full_test)
    file_path = os.path.join(save_path, "line_data")
    with open(file_path, 'wb') as f:
        pickle.dump(data_dict, f)

    # 读取保存的原始数据
    file_path = os.path.join(save_path, "line_data")
    with open(file_path, 'rb') as f:
        data_dict = pickle.load(f)

    # 训练集、测试集划分
    train_data, train_token, train_label, train_size, train_name = split_data(data_dict, "train")
    test_data, test_token, test_label, test_size, test_name = split_data(data_dict, "test")


    train_data, fea_transforms = prep_data(train_data)
    test_data, _ = prep_data(test_data, fea_transforms)
    # test_data, _ = prep_data(test_data)

    train_dict = {"data": train_data, "token": train_token, "label": train_label, "size": train_size}
    test_dict = {"data": test_data, "token": test_token, "label": test_label, "size": test_size}

    save_numpy_data(save_path, train_dict, "train")
    save_numpy_data(save_path, test_dict, "test")

    transforms_path = os.path.join(save_path, "transformation.pkl")
    with open(transforms_path, "wb") as file:
        pickle.dump(fea_transforms, file, pickle.HIGHEST_PROTOCOL)

    num_tokens = len(set(PHM_TOKEN.values()))
    feature_size = train_dict["data"].shape[-1]

    meta_information = dict()
    meta_information["num_tokens"] = num_tokens
    meta_information["feature_size"] = feature_size
    meta_information["full_life"] = PHM_LIFE
    meta_information["transforms_path"] = transforms_path
    meta_information["train_name"] = train_name
    meta_information["test_name"] = test_name

    meta_information_path = os.path.join(save_path, "meta_information.json")
    with open(meta_information_path, "w") as file:
        json.dump(meta_information, file)


def save_numpy_data(path, data_dict, suffix):
    for name, data in data_dict.items():
        np.save(os.path.join(path, f"{suffix}_{name}.npy"), data)


def train_test_split(data_dict, test_ratio=0.3):
    train_data, test_data = list(), list()
    train_label, test_label = list(), list()
    train_size, test_size = list(), list()
    data_token, data_name = list(), list()
    for name in TRAIN_BEARING + TEST_BEARING:
        seq_length = data_dict[name]["data"].shape[0]
        test_lengh = int(test_ratio * seq_length)
        train_data.append(data_dict[name]["data"][:-test_lengh])
        test_data.append(data_dict[name]["data"][test_lengh:])
        train_label.append(data_dict[name]["label"][:-test_lengh])
        test_label.append(data_dict[name]["label"][-test_lengh:])
        train_size.append(seq_length - test_lengh)
        test_size.append(test_lengh)
        data_token.append(data_dict[name]["token"])
        data_name.append(name)
    train_data = np.concatenate(train_data, axis=0)
    test_data = np.concatenate(test_data, axis=0)
    train_label = np.concatenate(train_label, axis=0)
    test_label = np.concatenate(test_label, axis=0)
    data_token = np.array(data_token)
    train_size = np.array(train_size)
    test_size = np.array(test_size)
    return train_data, test_data, train_label, test_label, train_size, test_size, data_token, data_name


def split_data(data_dict, select):
    if select == "train":
        bearings_name = TRAIN_BEARING
    else:
        bearings_name = TEST_BEARING
    bearings_data = list()
    bearings_size = list()
    bearings_label = list()
    bearings_token = list()
    for name in bearings_name:
        bearings_data.append(data_dict[name]["data"].copy())
        bearings_label.append(data_dict[name]["label"].copy())
        bearings_token.append(data_dict[name]["token"])
        bearings_size.append(data_dict[name]["data"].shape[0])
    bearings_data = np.concatenate(bearings_data, axis=0)
    bearings_token = np.array(bearings_token)
    bearings_label = np.concatenate(bearings_label, axis=0)
    bearings_size = np.array(bearings_size)
    return bearings_data, bearings_token, bearings_label, bearings_size, bearings_name


# 最大最小归一化
def prep_data(batch_data, transforms=None, keep_features=64):
    fea_size = batch_data.shape[1]

    if transforms is not None:
        assert len(transforms) == fea_size
        feature_transforms = transforms
    else:
        feature_transforms = [None for _ in range(fea_size)]

    # # 随机打乱数据
    # np.random.shuffle(batch_data)

    for i in range(fea_size):
        if transforms is None:
            # feature_transforms[i] = MinMaxScaler(feature_range=(0, 1))
            feature_transforms[i] = StandardScaler()
            batch_data[:, [i]] = feature_transforms[i].fit_transform(batch_data[:, [i]])
        else:
            batch_data[:, [i]] = feature_transforms[i].transform(batch_data[:, [i]])
    # if transforms is None:
    #     feature_transforms[-1] = KernelPCA(keep_features, kernel="rbf")
    #     batch_data = feature_transforms[-1].fit_transform(batch_data)
    # else:
    #     batch_data = feature_transforms[-1].transform(batch_data)

    return batch_data, feature_transforms


def load_phm2012(data_path, need_process, label_normalize=True, full_test=True):
    data_dict = {}
    set_folder_name = ["Learning_set"]
    if full_test:
        set_folder_name.append("Full_Test_Set")
    else:
        set_folder_name.append("Test_Set")
    for set_name in set_folder_name:
        dir_path = os.path.join(data_path, set_name)
        d = load_files(dir_path, need_process, label_normalize, full_test)
        data_dict.update(d)
    return data_dict


def load_files(dir_path, need_process, label_normalize, full_test):
    process_functions = [cal_fft_feature, cal_seq_feature, cal_time_feature, cal_wavelet_feature]
    # process_functions = [cal_fft_feature]
    bearings_folders = os.listdir(dir_path)
    bearings_data = dict()
    all_set = set(TRAIN_BEARING + TEST_BEARING)

    for bearing_name in bearings_folders:
        if bearing_name not in all_set:
            continue
        file_path = os.path.join(dir_path, bearing_name)
        bearing_data = []
        if not os.path.isdir(file_path):
            print(f"[warnning][{file_path} is not existed!]")
            continue
        file_names = os.listdir(file_path)
        file_names.sort()
        date = ["", ""]
        for i, file_name in enumerate(file_names):
            if not file_name.endswith("csv") or "acc" not in file_name:
                continue
            df = pd.read_csv(os.path.join(file_path, file_name), header=None)
            if df.shape[1] < 6:
                df = pd.read_csv(os.path.join(file_path, file_name), sep=";", header=None)
            if i == 0:
                date[0] = ":".join(str(int(x)) for x in df.loc[0, :2])
            else:
                date[1] = ":".join(str(int(x)) for x in df.loc[0, :2])
            data = np.array(df.loc[:, 4:6])
            data = data.transpose(1, 0)
            data = data[np.newaxis, :, :]
            bearing_data.append(data)
        if len(bearing_data) == 0:
            continue
        bearing_data = np.concatenate(bearing_data, axis=0)
        seq_size = bearing_data.shape[0]

        if need_process:
            bearing_data = np.concatenate([func(bearing_data) for func in process_functions], axis=1)
        else:
            bearing_data = bearing_data.reshape(seq_size, -1)

        # # 将所有为0的特征删除
        # bearing_data = bearing_data[:, ~(bearing_data == 0).all(axis=0)]

        rul = 0 if full_test else PHM_RUL[bearing_name]
        bearing_label = np.arange(seq_size)[::-1] * 10 + rul
        # bearing_label = np.arange(100)+ rul

        # # 划定分割点
        # point = int(Recession_point[bearing_name] / 10)
        # index_list = np.arange(point)
        # for index in index_list:
        #     bearing_label[index] = bearing_label[point]

        # 标签最大最小值归一化
        min_bearing_label = np.min(bearing_label)
        max_bearing_label = np.max(bearing_label)
        if label_normalize:
            bearing_label = (bearing_label - 0) / (max_bearing_label - 0)
            # bearing_label = bearing_label / PHM_LIFE
        bearing_token = PHM_TOKEN[bearing_name]

        # 将分割点之后的数据单独截取出来
        # bearing_data = bearing_data[point:]
        # bearing_label = bearing_label[point:]

        bearings_data[bearing_name] = {"data": bearing_data, "label": bearing_label, "token": bearing_token}
        print(
            f"finish load {bearing_name}, from {date[0]} to {date[1]}, the total length is {bearing_data.shape[0]}, first rul is {bearing_label[0]} ")


    return bearings_data


def cal_time_feature(data):
    fea_dict = OrderedDict()
    fea_dict['mean'] = np.mean(data, axis=2, keepdims=True)
    fea_dict['rms'] = np.sqrt(np.mean(data ** 2, axis=2, keepdims=True))
    # fea_dict['kur'] = np.sum((data-fea_dict['mean'].repeat(data.shape[2], axis=2))**4, axis=2) \
    #     / (np.var(data,axis=2) ** 2 * data.shape[2])
    # fea_dict['kur'] = fea_dict['kur'][:,:,np.newaxis]
    # fea_dict['skew'] = np.sum((data-fea_dict['mean'].repeat(data.shape[2],axis=2))**3,axis=2) \
    #     / (np.var(data,axis=2)**(3/2) * data.shape[2])
    # fea_dict['skew'] = fea_dict['skew'][:,:,np.newaxis]
    # fea_dict['p2p'] = np.max(data, axis=2, keepdims=True) - np.min(data, axis=2,keepdims=True)
    fea_dict['var'] = np.var(data, axis=2, keepdims=True)
    # fea_dict['cre'] = np.max(abs(data), axis=2, keepdims=True) / fea_dict['rms']
    # fea_dict['imp'] = np.max(abs(data), axis=2, keepdims=True) \
    #     / np.mean(abs(data), axis=2, keepdims=True)
    # fea_dict['mar'] = np.max(abs(data), axis=2, keepdims=True) \
    #     / (np.mean((abs(data)) ** 0.5, axis=2, keepdims=True))**2
    # fea_dict['sha'] = fea_dict['rms'] / np.mean(abs(data),axis=2,keepdims=True)
    # fea_dict['smr'] = (np.mean((abs(data))**0.5,axis=2,keepdims=True))**2
    # fea_dict['cle'] = fea_dict['p2p'] / fea_dict['smr']

    fea = np.concatenate(tuple(x for x in fea_dict.values()), axis=2)
    fea = fea.reshape(-1, fea.shape[1] * fea.shape[2])
    return fea


def cal_fft_feature(data, freq=2560):
    # hht变换
    # data_shape = data.shape
    # hilbert_data = np.zeros_like(data)
    # for i in range(data_shape[0]):
    #     for j in range(data_shape[1]):
    #         raw_data = data[i, j, :]
    #         eegRawHT = calculate_marginal_spectrum(raw_data)
    #         hilbert_data[i, j] = eegRawHT

    # 去除直流分量
    data = data - np.mean(data, axis=2, keepdims=True)
    sub_freq = int(freq / 5)
    freq_band = int(sub_freq / 2)
    fft_data = np.fft.fft(data, axis=2) / sub_freq

    fft_data = np.abs(fft_data) ** 2
    fea_list = []

    zero = int(fft_data.shape[2] / 2)

    for i in range(5):
        left = zero - i * freq_band
        right = zero + i * freq_band
        fea_list.append(np.sum(fft_data[:, :, left - freq_band: left], axis=2, keepdims=True) +
                        np.sum(fft_data[:, :, right: right + freq_band + 1], axis=2, keepdims=True))

    fea = np.concatenate(fea_list, axis=2)
    fea = fea.reshape(-1, fea.shape[1] * fea.shape[2])
    return fea


# 生成与序列长度等长的单减序列，每个轴承生成的序列长度不同，会导致在之后的
# 标准后，短的序列中特征无法从归一到0-1。
def cal_seq_feature(data):
    n = data.shape[0]
    x = np.arange(n)[::-1].reshape(n, 1)
    return x


def cal_wavelet_feature(data, freq=2560, num_levels=4, wavelet_function="db4", keep_size=1):
    wp = pywt.WaveletPacket(data, wavelet=wavelet_function, maxlevel=num_levels, axis=-1)

    packet_names = [node.path for node in wp.get_level(num_levels, "natural")]
    features = []
    num_samples = data.shape[-1]
    f = np.linspace(0.0, (freq / 2.0), num_samples // 2)
    for p in packet_names:
        new_wp = pywt.WaveletPacket(data=None, wavelet="db4", maxlevel=num_levels, axis=-1)
        new_wp[p] = wp[p].data
        reconstructed_signal = new_wp.reconstruct(update=False)
        reconstructed_feature = cal_fft_feature(reconstructed_signal, freq)
        features += [reconstructed_feature]
        # freq_values = np.fft.fft(reconstructed_signal, axis=-1)
        # freq_values = 2.0 / num_samples * np.abs(freq_values[...,0: (num_samples // 2)])
        # z = abs(freq_values)
        # maximal_idx = np.argpartition(z, keep_size, axis=-1)[..., :keep_size]
        # high_amp = np.take_along_axis(z, maximal_idx, axis=-1)
        # high_freq = f[maximal_idx]
        # features += [high_amp, high_freq]
    features = np.concatenate(features, axis=-1).reshape(data.shape[0], -1)
    return features


import numpy as np
from scipy.signal import hilbert
from pyhht import EMD
from scipy.stats import pearsonr


def calculate_marginal_spectrum(data):
    # Step 1: Perform EMD on the signal
    decomposer = EMD(data)
    imfs = decomposer.decompose()

    marginal_spectrum = np.zeros_like(data)

    # Step 2: Apply Hilbert transform on each IMF to get instantaneous frequency
    for imf in imfs:
        analytic_signal = hilbert(imf)
        instantaneous_phase = np.unwrap(np.angle(analytic_signal))
        instantaneous_frequency = (np.diff(instantaneous_phase) / (2.0 * np.pi) * 2560)

        # Step 3: Accumulate instantaneous frequencies of all IMFs to get the marginal spectrum
        for i in range(len(instantaneous_frequency)):
            freq_index = int(instantaneous_frequency[i])
            marginal_spectrum[freq_index] += abs(analytic_signal[i])

    return marginal_spectrum


def Pearson_correlation(batch_data):
    first_row = batch_data[0, :]
    correlations = []

    # Iterate over each row in batch_data
    for i in range(batch_data.shape[0]):
        # Calculate the Pearson correlation coefficient between the first row and the current row
        corr, _ = pearsonr(first_row, batch_data[i, :])
        # Append the correlation coefficient to the list
        correlations.append(corr)

    return correlations


if __name__ == "__main__":
    print("# Loading the data:")

    parser = argparse.ArgumentParser()
    parser.add_argument("--data-path", type=str, dest="data_path", help="phm data path",
                        default="D:/dataset/ieee-phm-2012-data-challenge-dataset/")
    parser.add_argument("--prep", dest="prep", action="store_true", help="run preprocess for data or not.")
    parser.add_argument("--no-prep", dest="prep", action="store_false")
    parser.set_defaults(prep=True)

    parser.add_argument("--test", dest="full_test", action="store_false", help="")
    parser.add_argument("--full-test", dest="full_test", action="store_true", help="")
    parser.set_defaults(full_test=True)
    parser.add_argument("--save-path", type=str, dest="save_path", help="meta information save path",
                        default="./data/phm2012")

    args = parser.parse_args()
    data_path = args.data_path
    save_path = args.save_path
    preprocess = args.prep
    full_test = args.full_test

    main(data_path, save_path, preprocess, full_test)
